import socket
import threading
from cryptography.fernet import Fernet

# 从文件中加载密钥
def load_key():
    with open("secret.key", "rb") as key_file:
        return key_file.read()
key = load_key()
cipher = Fernet(key)

# 服务器主机和端口
HOST = '0.0.0.0'  # 监听所有地址
PORT = 12345  # 你可以修改这个端口

# 创建一个 TCP/IP 套接字
server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server_socket.bind((HOST, PORT))
server_socket.listen(5)

clients = []  # 保存所有客户端连接
nicknames = []  # 保存所有客户端昵称

# 向所有客户端广播消息
def broadcast(message):
    for client in clients:
        client.send(message)

# 处理客户端消息
def handle_client(client):
    while True:
        try:
            # 接收客户端加密的消息
            encrypted_message = client.recv(1024)
            message = cipher.decrypt(encrypted_message).decode('utf-8')  # 解密消息并转换为UTF-8格式
            broadcast(encrypted_message)  # 将加密的消息广播
        except:
            # 如果客户端断开连接
            index = clients.index(client)
            clients.remove(client)
            client.close()
            nickname = nicknames[index]
            broadcast(cipher.encrypt(f'{nickname} has left the chat.'.encode('utf-8')))
            nicknames.remove(nickname)
            break

# 接收客户端连接
def receive():
    print("Server is running and listening for connections...")
    while True:
        # 接收客户端连接
        client, address = server_socket.accept()
        print(f"Connected with {str(address)}")

        # 请求并存储客户端的昵称
        client.send(cipher.encrypt("NICK".encode('utf-8')))
        nickname = cipher.decrypt(client.recv(1024)).decode('utf-8')
        nicknames.append(nickname)
        clients.append(client)

        print(f"Nickname of the client is {nickname}!")
        broadcast(cipher.encrypt(f"{nickname} joined the chat!".encode('utf-8')))
        client.send(cipher.encrypt("You are now connected to the chat!".encode('utf-8')))

        # 为每个客户端创建一个线程处理消息
        thread = threading.Thread(target=handle_client, args=(client,))
        thread.start()

# 启动服务器
if __name__ == "__main__":
    receive()
